feat(usage): always-on per-query usage tracking#614
Conversation
Add a dedicated, provider-agnostic usage-tracking layer that records every query onto the central Organizations graph, independent of the optional LLM conversational-memory feature. Background: the only pre-existing per-query record was a side effect of the opt-in `use_memory` feature, which is also gated to OpenAI/Azure providers and lazily created. As a result the vast majority of users had no recorded query activity, and every recorded query was against the demo DB — usage counts were unusable for measuring adoption. This change: - Adds api/core/usage_tracking.py with record_query_usage_background(), a fire-and-forget recorder mirroring save_memory_background (background task, task-sink support for the SDK, errors logged-not-raised). - Maintains denormalized counters on the User node (query_count/success_count/ error_count/last_active/first_query_at) plus a per-query (:UsageEvent) node linked (User)-[:PERFORMED]-> for time-series/per-DB/success-rate analytics. - Hooks the recorder into run_query and run_confirmed at the completion point, OUTSIDE the use_memory/provider gate, so it runs on every query. - Uses MATCH (not MERGE) on User so an unknown email is a silent no-op rather than creating a phantom user from the query path. Forward-only: historical usage cannot be backfilled. Tests: tests/test_usage_tracking.py (write content, demo flag, ungated design, invalid user_id no-op, swallowed write failure). Cypher validated against a live FalkorDB engine. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Completed Working on "Code Review"✅ Review publishing completed successfully. Posted comments from all chunks and submitted final review: COMMENT with 1 total comment across 1 file. ✅ Workflow completed successfully. |
Dependency Review✅ No vulnerabilities or license issues or OpenSSF Scorecard issues found.Scanned FilesNone |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 WalkthroughWalkthroughAdds a configurable Organizations graph name, routes auth and token queries through it, records per-query usage in a background task, wires that tracking into streaming query flows, and adds unit tests. ChangesPer-query usage tracking
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Review Summary
Findings by severity:
- MAJOR: 1
- BLOCKER/CRITICAL/MINOR/SUGGESTION/PRAISE: 0
Affected files:
api/core/text2sql.py
Key theme:
- Usage tracking coverage is incomplete on early-exit paths, which can undercount real query activity and conflict with the always-on tracking goal.
Actionable next steps:
- Ensure usage recording runs on all terminal paths in
run_queryandrun_confirmed(including early returns and exceptions), ideally via a singlefinally-style completion hook. - Add/extend tests to cover early-exit/error paths to verify exactly-once usage recording behavior.
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
api/core/usage_tracking.py (1)
40-53: 🚀 Performance & Scalability | 🔵 TrivialPlan for unbounded
UsageEventgrowth.Every query
CREATEs a new(:UsageEvent)node on the sharedOrganizationsgraph with no retention bound. Over time this graph grows without limit, and time-series/per-DB analytics queries will scan ever-larger node sets. Consider an index onUsageEvent(timestamp)(and/orgraph_id) for read efficiency, plus a retention/rollup strategy.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@api/core/usage_tracking.py` around lines 40 - 53, The `record_usage` Cypher in `_RECORD_USAGE_CYPHER` creates an unbounded `UsageEvent` node for every query, so add a retention or rollup strategy to prevent indefinite growth while keeping analytics usable. Update the usage-tracking flow around `UsageEvent` creation to either prune older events, aggregate them into summary records, or both, and add indexes on `UsageEvent.timestamp` and/or `graph_id` to keep reads efficient. Keep the `MATCH (u:User ...)` update behavior intact while making the event-storage path bounded.tests/test_usage_tracking.py (1)
143-150: 📐 Maintainability & Code Quality | 🔵 TrivialConsider a test for the no-sink path.
All
TestRecordQueryUsagecases pass an explicittask_sink, so thesink is Nonebranch (the one used by the server path viabackground_tasks_var) is never exercised. A test covering the no-sink scheduling path would guard against the task-GC concern raised inusage_tracking.py.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_usage_tracking.py` around lines 143 - 150, Add a test that exercises the no-sink scheduling path in the usage tracking flow, since all current TestRecordQueryUsage cases pass an explicit task_sink. Cover the branch where record_query_usage_background is invoked with sink/task_sink absent and the server-style background_tasks_var path is used, and assert the scheduled task still runs or is retained correctly. Use the existing TestRecordQueryUsage or TestUngatedDesign setup and reference record_query_usage_background and background_tasks_var to place the new case alongside the current usage-tracking tests.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@api/core/text2sql.py`:
- Around line 612-616: The usage tracking call in text2sql is currently placed
too late, so early-return paths like off-topic, non-translatable, demo-block,
and confirmation exits bypass it and are never counted. Move
record_query_usage_background(...) earlier in the query flow in the text2sql
handling logic so it runs before any of those returns, while keeping its
existing success flag and db arguments unchanged.
In `@api/core/usage_tracking.py`:
- Around line 86-89: The usage_event INFO log in usage_tracking.py is emitting
user-derived email and graph_id values without sanitization, which can allow log
forging via CR/LF. Update the logging path around the usage_event call to
neutralize or escape these fields before passing them to logging.info, and
consider reducing or removing the email PII from this high-volume INFO log if it
is not required. Use the usage_event logging statement as the place to apply the
fix.
---
Nitpick comments:
In `@api/core/usage_tracking.py`:
- Around line 40-53: The `record_usage` Cypher in `_RECORD_USAGE_CYPHER` creates
an unbounded `UsageEvent` node for every query, so add a retention or rollup
strategy to prevent indefinite growth while keeping analytics usable. Update the
usage-tracking flow around `UsageEvent` creation to either prune older events,
aggregate them into summary records, or both, and add indexes on
`UsageEvent.timestamp` and/or `graph_id` to keep reads efficient. Keep the
`MATCH (u:User ...)` update behavior intact while making the event-storage path
bounded.
In `@tests/test_usage_tracking.py`:
- Around line 143-150: Add a test that exercises the no-sink scheduling path in
the usage tracking flow, since all current TestRecordQueryUsage cases pass an
explicit task_sink. Cover the branch where record_query_usage_background is
invoked with sink/task_sink absent and the server-style background_tasks_var
path is used, and assert the scheduled task still runs or is retained correctly.
Use the existing TestRecordQueryUsage or TestUngatedDesign setup and reference
record_query_usage_background and background_tasks_var to place the new case
alongside the current usage-tracking tests.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 0e4abbeb-6450-4184-9a14-ac8c450e348e
📒 Files selected for processing (3)
api/core/text2sql.pyapi/core/usage_tracking.pytests/test_usage_tracking.py
There was a problem hiding this comment.
Pull request overview
Adds an always-on, provider-agnostic usage tracking mechanism to record per-query activity in the central Organizations graph, independent of the optional conversational-memory feature, to enable reliable adoption/usage analytics.
Changes:
- Introduces
api/core/usage_tracking.pywith a fire-and-forget recorder that updates denormalized counters onUserand appendsUsageEventnodes. - Hooks usage recording into both
run_queryandrun_confirmedcompletion points inapi/core/text2sql.py. - Adds unit tests validating decoding, write parameters, demo-flagging, failure swallowing, and ungated design.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| api/core/usage_tracking.py | New background usage recorder + Cypher write into Organizations graph. |
| api/core/text2sql.py | Integrates usage tracking calls into the query and confirmed-destructive flows. |
| tests/test_usage_tracking.py | Adds unit tests for the new tracking module behavior and API shape. |
|
🚅 Deployed to the QueryWeaver-pr-614 environment in queryweaver
|
- text2sql: record usage on every terminal path of run_query/run_confirmed (off-topic, not-translatable, demo-blocked, cancelled, no-loader), via a local _track_usage helper. The destructive-confirmation prompt is the deliberate exception — run_confirmed records that query's outcome, so tracking it twice would double-count. - usage_tracking: omit email (PII) from the per-query log line and strip CR/LF from the user-influenced graph_id (CodeQL log-injection). - usage_tracking: base64 decode with validate=True + email-shape check so a malformed user_id is skipped instead of triggering a phantom write. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Two design fixes raised in review: 1. Single source of truth for the central graph name. Add config.ORGANIZATIONS_GRAPH (env var ORGANIZATIONS_GRAPH, default "Organizations") and use it everywhere the graph was hardcoded — auth/user_management, routes/auth, routes/tokens, and usage tracking. 2. Don't ship hosted-app telemetry in the PyPI SDK. Move usage_tracking from api/core (shipped) to api/routes (excluded from the wheel) and invoke it from the route layer (_serialize_pipeline) instead of inside run_query/run_confirmed in api/core/text2sql.py. The route records exactly once from the final QueryResult and skips the destructive- confirmation prompt (the /confirm call records that query). text2sql.py is back to its pristine, tracking-free state. Verified: SDK wheel no longer contains usage_tracking and text2sql has no tracking refs; unit suite 138 passed, pylint 10/10. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
api/routes/auth.py (1)
137-174: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winInitialize
safe_emailbefore thetry.If
db.select_graph(ORGANIZATIONS_GRAPH)fails, theexceptblock logssafe_emailbefore it has been assigned, so the original error gets replaced byUnboundLocalError.Suggested fix
async def _set_mail_hash(email: str, password_hash: str) -> bool: """Set email hash for the user in the database.""" + safe_email = _sanitize_for_log(email) try: organizations_graph = db.select_graph(ORGANIZATIONS_GRAPH) - - # Sanitize inputs for logging - safe_email = _sanitize_for_log(email) # Create new email identity and user🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@api/routes/auth.py` around lines 137 - 174, Initialize safe_email before entering the try block in _set_mail_hash so it is always available in the except path. The current logging in the exception handler references safe_email after db.select_graph(ORGANIZATIONS_GRAPH) can fail, which can mask the original error with an UnboundLocalError; move the _sanitize_for_log(email) assignment above the try and keep the existing logging/HTTPException flow unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@api/config.py`:
- Around line 17-20: The ORGANIZATIONS_GRAPH config currently preserves an empty
string, which causes downstream db.select_graph(ORGANIZATIONS_GRAPH) calls to
use an invalid graph name. Update the configuration in config.py so the
Organizations default is used when the environment variable is missing or blank,
and keep the change localized to the ORGANIZATIONS_GRAPH constant and its
surrounding fallback logic.
In `@api/routes/graphs.py`:
- Around line 57-60: Early exceptions in the graph query flow can return before
the _Final sentinel is set, so usage tracking is skipped entirely. Update the
route-level error handling in graphs.py and the query handlers around
run_query() and run_confirmed() so record_query_usage_background is invoked for
any failed request, including pre-result failures like get_db_description()
awaits; use the existing user_id, namespaced, and success=False path even when
the exception is caught before final is produced.
- Line 37: The _serialize_pipeline helper is still missing explicit typing, so
update its signature to annotate gen as AsyncGenerator[dict | _Final, None],
user_id as str, namespaced as str, and declare the return type as
AsyncGenerator[str, None]. Make the change directly on _serialize_pipeline so it
matches the repository typing rule and keeps pylint satisfied.
---
Outside diff comments:
In `@api/routes/auth.py`:
- Around line 137-174: Initialize safe_email before entering the try block in
_set_mail_hash so it is always available in the except path. The current logging
in the exception handler references safe_email after
db.select_graph(ORGANIZATIONS_GRAPH) can fail, which can mask the original error
with an UnboundLocalError; move the _sanitize_for_log(email) assignment above
the try and keep the existing logging/HTTPException flow unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: aa885f91-0cc5-41ce-9dd3-832feb5d8760
📒 Files selected for processing (8)
.env.exampleapi/auth/user_management.pyapi/config.pyapi/routes/auth.pyapi/routes/graphs.pyapi/routes/tokens.pyapi/routes/usage_tracking.pytests/test_usage_tracking.py
✅ Files skipped from review due to trivial changes (1)
- .env.example
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/test_usage_tracking.py
…ack crashes
- usage_tracking: the per-query log line logged the namespaced graph_id
({base64(email)}_{db}); base64 email is reversible, so log a short SHA-256
hash instead — keeps identity out of logs and neutralizes log-injection.
- config: ORGANIZATIONS_GRAPH falls back to "Organizations" when the env var
is empty (not just unset), so a blank value can't target an empty graph.
- routes/graphs: record success=False when run_query/run_confirmed raise
before the _Final sentinel (e.g. get_db_description), so pipeline crashes
are still counted instead of silently bypassing tracking.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Address the latest review pass: - routes/graphs: success = is_valid AND no error. error_message-None alone counted off-topic / not-SQL-translatable results (is_valid=False, no error) as successes, inflating success_count. - tests: assert against usage_tracking.ORGANIZATIONS_GRAPH instead of the hardcoded "Organizations" so the suite passes when the env var is set. - usage_tracking: reword the task_sink docstring — it no longer ships in the SDK, so drop the QueryWeaver.close() reference for a generic description. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Why
While analyzing production usage we found that 235 of 240 users showed zero recorded queries — not because they weren't using the product, but because QueryWeaver has no usage-tracking system. The only per-query record was a side effect of the optional LLM conversational-memory feature (
api/memory/graphiti_tool.py), which is:use_memory(defaultFalsefor the SDK),-memorygraph only exists after a qualifying query, and failures are swallowed.Consequently every recorded query in production was against the demo DB (
DEMO_CRM) — not a single own-database query was ever tracked. Usage counts were unusable for measuring adoption.What
Adds a dedicated, always-on, provider-agnostic usage-tracking layer that records every query, independent of the memory feature.
api/core/usage_tracking.py—record_query_usage_background(...), a fire-and-forget recorder mirroringsave_memory_background(background task, SDK task-sink support, errors logged-not-raised).Organizationsgraph (alongsideUser/Identity/Token):Usernode:query_count,success_count,error_count,last_active,first_query_at;(:UsageEvent {graph_id, is_demo, success, timestamp})linked(User)-[:PERFORMED]->for time-series, per-DB and success-rate analytics (is_demofixes the "all data isDEMO_CRM" blind spot).run_queryandrun_confirmedat the completion point, outside theuse_memory/provider gate, so it runs on every query. No route or signature changes.MATCH(notMERGE) onUser, so an unknown email is a silent no-op rather than a phantom user.Notes / scope
Episodic/Community/Saganever created, demo-only writes) and a read/admin API + frontend usage view.Testing
tests/test_usage_tracking.py(9 tests): write content for success/error, demo flag, ungated design (recorder has nouse_memory/providerparams), invaliduser_idno-op, swallowed write failure. Marked@pytest.mark.unit.UsageEventnodes are created, and an unknown email is a confirmed no-op (no phantom user). Production was not written to.🤖 Generated with Claude Code
Summary by CodeRabbit
Organizations).